# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import logging
import math
import sys
import warnings

import MySQLdb
import MySQLdb.cursors
import MySQLdb.converters

from python.base.init import DEFINE_bool, FLAGS
from python.database.base.db_user import DbUser

DEFINE_bool('dry', False, 'dry run')

DEFINE_bool('db_manager_error_on_warning', False, '')


# DbManager will provide a low-level interface of database access.
# It's role is to assemble a query that is frequently used.
# Since this part is extremely performance-sensitive, this tries to
# optimize each query to function properly.
#
# Also, a guest quant user must call read-only db manager:
# QuantDbManager()
#
# An administrator should do this:
# AdminDbManager()
class DbManager(object):
  @staticmethod
  def get_default_typing():
    field_typing = MySQLdb.converters.conversions
    field_typing[MySQLdb.converters.FloatType] = (lambda s, d: "NULL" if math.isnan(s) else str(s))
    field_typing[MySQLdb.FIELD_TYPE.DATETIME] = str
    field_typing[MySQLdb.FIELD_TYPE.DATE] = str
    field_typing[MySQLdb.FIELD_TYPE.TIME] = str
    return field_typing

  def __init__(self, db_user, charset="utf8"):
    self._db_user = db_user
    self.connect(charset)

  def connect(self, charset):
    self._fetch_how = 1
    self._db = MySQLdb.connect(
        conv=self.get_default_typing(),
        host=self._db_user['database_url'],
        user=self._db_user['user'],
        passwd=self._db_user['password'],
        db='' if self._db_user['database_name'] == '*' else self._db_user['database_name'],
        charset=charset)
    self._reset_cursor()
    self.initialized = False
    if FLAGS.db_manager_error_on_warning:
      warnings.filterwarnings('error', category=MySQLdb.Warning)

    if FLAGS.dry:
      logging.info(self._db)

    self._query("SHOW VARIABLES LIKE 'max_allowed_packet'")
    self._max_allowed_packet = int(self._fetch()[0]["Value"])
    self.initialized = True

  def literal(self, value):
    return self._db.literal(value)

  def set_character_set(self, charset):
    self._query("SET CHARACTER SET %s", (charset,))

  def _create_tuple(self, strings):
    return ", ".join(strings)

  def _create_value_tuple(self, values):
    return ",".join([self.literal(value) for value in values])

  def _create_where_tuple(self, where_dict, like):
    if not like:
      query_str = " and ".join(
          ["(`%s`=%s)" % (key, self.literal(value)) for key, value in where_dict.iteritems()])
    else:
      query_str = " and ".join(
          ["(`%s` like %s)" % (key, self.literal(value)) for key, value in where_dict.iteritems()])
    query_str = query_str.replace("`=NULL", "` IS NULL").replace("` like NULL", "` IS NULL")
    return query_str

  def _create_where(self, where_dicts, like):
    if len(where_dicts) > 0:
      where_strs = [self._create_where_tuple(where_dict, like) for where_dict in where_dicts]
      if len(where_strs) > 1:
        return "WHERE (%s)" % " or ".join(["(" + where_str + ")" for where_str in where_strs])
      else:
        return "WHERE %s" % where_strs[0]
    else:
      return ""

  def _create_set(self, set_dict):
    strs = []
    for key, value in set_dict.iteritems():
      strs.append("`%s`=%s" % (key, self.literal(value)))
    return ",".join(strs)

  def _delete_from_query(self, table_name, where_dicts, like):
    sql_command = "DELETE FROM %s %s" % (table_name, self._create_where(where_dicts, like))
    return sql_command

  def _select_from_query(self, table_name, value_select, where_dicts, like, optional_str):
    sql_command = "SELECT %s FROM %s %s %s" % (self._create_tuple(value_select),
                                               table_name,
                                               self._create_where(where_dicts, like),
                                               optional_str)
    return sql_command

  def _update_query(self, table_name, where_dict, set_dict, like):
    sql_command = "UPDATE %s SET %s %s" % (
        table_name, self._create_set(set_dict), self._create_where([where_dict], like))
    return sql_command

  def __drop_table(self, table_name):
    self._query("DROP TABLE %s" % table_name)

  def break_args(self, tuples):
    arg_packet_sum = 0
    iter_tuples = []
    for a_tuple in tuples:
      arg_packet_sum += sum([sys.getsizeof(value) for value in a_tuple])
      iter_tuples.append(a_tuple)
      if arg_packet_sum * 2 > self._max_allowed_packet:
        yield iter_tuples
        arg_packet_sum = 0
        iter_tuples = []
    yield iter_tuples

  def commit(self):
    self._db.commit()

  def rollback(self):
    self._db.rollback()

  def start_transaction(self):
    self._query("START TRANSACTION")

  def use_database(self, database_name):
    self._query("USE %s" % database_name)

  def _query(self, query, args=None, many=False):
    self._reset_cursor()
    if FLAGS.dry:
      logging.info(query)
      logging.info(args)
    else:
      try:
        if many:
          for args_executed in self.break_args(args):
            self._cursor.executemany(query, args_executed)
        else:
          self._cursor.execute(query, args)
      except MySQLdb.Error:
        print query
        logging.info(query)
        if hasattr(self._cursor, '_last_executed'):
          print(self._cursor._last_executed)
          logging.info(self._cursor._last_executed)
        raise

  def _reset_cursor(self):
    if self._fetch_how == 0:
      self._cursor = self._db.cursor(MySQLdb.cursors.Cursor)
    else:
      assert self._fetch_how == 1
      self._cursor = self._db.cursor(MySQLdb.cursors.DictCursor)

  def get_lastrowid(self):
    return self._cursor.lastrowid

  def get_affected_rows(self):
    return self._cursor.rowcount

  def set_fetch_how(self, fetch_how):
    assert fetch_how in [0, 1]
    self._fetch_how = fetch_how

  def _fetch(self):
    if FLAGS.dry:
      pass
    else:
      return self._cursor.fetchall()

  def does_table_exist(self, table_name, database_name=None):
    if database_name is not None:
      self._query("SHOW TABLES FROM %s WHERE tables_in_%s LIKE %s" %
                  (database_name, database_name, self.literal(table_name)))
    else:
      self._query("SHOW TABLES LIKE %s", (table_name,))
    result = self._fetch()
    return len(result) > 0

  def create_table_if_nonexists(self, table_info, temporary=False):
    if not self.does_table_exist(table_info.get_table_name()):
      self.create_table(table_info, temporary)

  def create_table(self, table_info, temporary=False):
    query_to_create_table = table_info.get_create_table_query(temporary)
    self._query(query_to_create_table)

  def count_table(self, table_name, where_dict):
    sql_command = "SELECT COUNT(*) FROM %s %s" % (table_name,
                                                  self._create_where(where_dict, like=False))
    self._query(sql_command)
    return int(self._fetch()[0].values()[0])

  def drop_table_if_empty(self, table_name):
    element_count = self.count_table(table_name, {})
    if element_count == 0:
      logging.info("%s dropping table.." % table_name)
      self.__drop_table(table_name)
    else:
      logging.info("%s table not empty" % table_name)

  def _combine_insert_query(self, base_command, value_tuples):
    query = base_command
    if len(value_tuples) > 0:
      sub_query = "),(".join(value_tuples)
      query += ",(" + sub_query + ")"
    return query

  def query_on_kv_pair(self,
                       query_name,
                       table_name,
                       keys,
                       values,
                       appended_query=None,
                       autocommit=True):
    if appended_query is None:
      appended_query = ""
    self._query("%s %s (%s) VALUES (%s) %s " % (query_name,
                                                table_name,
                                                self._create_tuple(keys),
                                                self._create_tuple(["%s" for key in keys]),
                                                appended_query), [tuple(value) for value in values],
                many=True)
    if autocommit:
      self.commit()

  def insert_on_duplicate_update_data_kv_pair(self,
                                              table_name,
                                              set_keys,
                                              set_values,
                                              autocommit=True):
    on_duplicate_query_str = self._on_duplicate_query(
        {set_key: set_value for set_key, set_value in zip(set_keys, set_values[0])})
    self.query_on_kv_pair("INSERT INTO",
                          table_name,
                          set_keys,
                          set_values,
                          " ON DUPLICATE KEY UPDATE %s" % on_duplicate_query_str,
                          autocommit=autocommit)

  def _on_duplicate_query_notnull(self, set_dict):
    key_queries = [
        "`%s` = COALESCE(VALUES(`%s`), `%s`)" % (key, key, key) for key in set_dict.keys()
    ]
    return ",".join(key_queries)

  def insert_on_duplicate_update_notnull_kv(self,
                                            table_name,
                                            set_keys,
                                            set_values,
                                            autocommit=True,
                                            update_ignore_columns=None):
    if update_ignore_columns is None:
      update_ignore_columns = []
    on_duplicate_query_str = self._on_duplicate_query_notnull({
        set_key: set_value for set_key,
        set_value in zip(set_keys, set_values[0]) if set_key not in update_ignore_columns
    })
    self.query_on_kv_pair("INSERT INTO",
                          table_name,
                          set_keys,
                          set_values,
                          " ON DUPLICATE KEY UPDATE %s" % on_duplicate_query_str,
                          autocommit=autocommit)

  def query_data(self, query_name, table_name, dicts, autocommit=False):
    self.query_on_kv_pair(query_name,
                          table_name, [] if len(dicts) == 0 else dicts[0].keys(),
                          [row.values() for row in dicts],
                          autocommit=autocommit)

  def insert_data(self, *args, **kwargs):
    self.query_data("INSERT INTO", *args, **kwargs)

  def insert_on_duplicate_update_data(self, table_name, dicts, autocommit=True):
    self.insert_on_duplicate_update_data_kv_pair(table_name,
                                                 [] if len(dicts) == 0 else dicts[0].keys(),
                                                 [row.values() for row in dicts],
                                                 autocommit=autocommit)

  @staticmethod
  def _chunks(l, n):
    n = max(1, n)
    return [l[i:i + n] for i in range(0, len(l), n)]

  def delete_data(self, table_name, delete_dicts, like, autocommit=True):
    affected_rows = 0
    if len(delete_dicts) == 0:
      return 0
    elif len(delete_dicts) < 10:
      sql_command_indiv = self._delete_from_query(table_name, delete_dicts, like)
      self._query(sql_command_indiv)
      affected_rows += self.get_affected_rows()
    else:
      for chunk_delete_dicts in self._chunks(delete_dicts, 10):
        sql_command_indiv = self._delete_from_query(table_name, chunk_delete_dicts, like)
        self._query(sql_command_indiv)
        affected_rows += self.get_affected_rows()
    if autocommit:
      self.commit()
    return affected_rows

  def select_data_raw_query(self, query):
    self._query(query)
    return self._fetch()

  def execute_raw_query(self, query, args=None, autocommit=True):
    self._query(query, args, False)
    if autocommit:
      self.commit()

  def select_execute_raw_query(self, query, args=None, autocommit=False):
    self._query(query, args, False)
    if autocommit:
      self.commit()
    return self._fetch()

  def fetch(self):
    return self._fetch()

  def select_data(self, table_name, columns_select, select_dicts, like, optional_str):
    if len(select_dicts) < 10:
      return self.__select_data_slow(table_name, columns_select, select_dicts, like, optional_str)

    temp_table_name = table_name.replace(".", "_")

    pkey_str = ",".join(['`%s`' % key for key in select_dicts[0].keys()])

    sql_command = '''
CREATE TEMPORARY TABLE
temporary_tables.%s_temp (PRIMARY KEY (%s)) Engine=MEMORY
  SELECT %s FROM %s WHERE FALSE''' % (temp_table_name, pkey_str, pkey_str, table_name)
    self._query(sql_command)

    results = []
    self.query_data("INSERT INTO",
                    "temporary_tables.%s_temp" % temp_table_name,
                    select_dicts,
                    autocommit=False)

    if len(select_dicts[0].keys()) == 1:
      sql_command = '''
SELECT %s from %s where exists ( select 1 from temporary_tables.%s_temp where
temporary_tables.%s_temp.%s = %s.%s limit 1) %s''' % (self._create_tuple(columns_select),
                                                      table_name,
                                                      temp_table_name,
                                                      temp_table_name,
                                                      pkey_str.replace("`", ""),
                                                      table_name,
                                                      pkey_str.replace("`", ""),
                                                      optional_str)
    else:
      sql_command = '''
SELECT %s from %s NATURAL JOIN temporary_tables.%s_temp where true %s;
''' % (self._create_tuple(columns_select), table_name, temp_table_name, optional_str)

    self._query(sql_command)
    results += self._fetch()
    self._query("DROP TEMPORARY TABLE temporary_tables.%s_temp" % temp_table_name)
    return results

  def __select_data_slow(self, table_name, columns_select, select_dicts, like, optional_str):
    results = []
    if len(select_dicts) > 0:
      for chunk_select_dicts in self._chunks(select_dicts, 5):
        sql_command = self._select_from_query(table_name,
                                              columns_select,
                                              chunk_select_dicts,
                                              like,
                                              optional_str)
        self._query(sql_command)
        results += self._fetch()
    else:
      sql_command = self._select_from_query(table_name,
                                            columns_select,
                                            select_dicts,
                                            like,
                                            optional_str)
      self._query(sql_command)
      results += self._fetch()
    return results

  def update_data(self, table_name, where_dicts, set_dicts, like, autocommit=True):
    assert len(where_dicts) == len(set_dicts)
    for where_dict, set_dict in zip(where_dicts, set_dicts):
      sql_command = self._update_query(table_name, where_dict, set_dict, like)
      self._query(sql_command)
    if autocommit:
      self.commit()

  def _on_duplicate_query(self, set_dict):
    key_queries = ["`%s` = VALUES(`%s`)" % (key, key) for key in set_dict.keys()]
    return ",".join(key_queries)

  def extract_columns_from_schema(self, table_name, schema):
    self._query("CREATE TEMPORARY TABLE _%s_temp %s" % (table_name, schema))
    self._query("SHOW INDEX from _%s_temp" % table_name)
    index_keys = self.show_index_to_index_keys(self._fetch())
    self._query("DESC _%s_temp" % table_name)
    column_list = self.desc_to_column_list(self._fetch(), index_keys)
    self._query("DROP TABLE _%s_temp" % table_name)
    return column_list

  def show_index_to_index_keys(self, fetch_results):
    key_list = []
    for fetch_result in fetch_results:
      if fetch_result['Key_name'].lower() != "primary":
        key_list.append(fetch_result['Column_name'])
    return key_list

  def desc_to_column_list(self, fetch_results, index_keys):
    column_list = []
    for fetch_result in fetch_results:
      column_desc = {
          'column_name': fetch_result['Field'],
          'column_type': fetch_result['Type'],
          'primary_key': fetch_result['Key'].lower() == "pri",
          'nullable': (fetch_result['Null'].lower() == "yes")
      }
      if column_desc['column_name'] in index_keys:
        column_desc['index'] = True
      column_list.append(column_desc)
    return column_list

  def get_column_list(self, table_name):
    self._query("DESC %s" % table_name)
    column_list = self.desc_to_column_list(self._fetch(), [])
    return column_list

  def create_user(self, a_db_user):
    sql_command = "CREATE USER '%(user)s'@'%%' IDENTIFIED BY '%(password)s'" % (a_db_user)
    self._query(sql_command)

  def drop_user(self, a_db_user):
    sql_command = "DROP USER '%(user)s'@'%%'" % (a_db_user)
    self._query(sql_command)

  def does_user_exist(self, user_name):
    self._query("select User,Host from mysql.user where user like '%s';" % user_name)
    return len(self._fetch()) > 0

  def create_all_users(self, json_config_file):
    user_settings = DbUser.load_users_from_file(json_config_file)
    if not self.does_database_exist("temporary_tables"):
      self.create_database("temporary_tables")
    for name, setting in user_settings.iteritems():
      if (setting["database_name"] != "*"
          and not self.does_database_exist(setting["database_name"])):
        self.create_database(setting["database_name"])
      if not self.does_user_exist(setting["user"]):
        self.create_user(setting)
      if setting["privilege"] == "admin":
        self._query("GRANT ALL PRIVILEGES ON %(database_name)s.* TO '%(user)s'@'%%'" % setting)
      elif setting["privilege"] == "readonly":
        self._query("GRANT SELECT ON %(database_name)s.* TO '%(user)s'@'%%'" % setting)
        self._query("""GRANT SELECT, INSERT, UPDATE, DELETE, CREATE, DROP,
            INDEX, ALTER, CREATE TEMPORARY TABLES ON temporary_tables.*
            TO '%(user)s'@'%%'""" % setting)
      else:
        raise ValueError(setting["privilege"])

  def does_database_exist(self, database_name):
    sql_command = "SHOW DATABASES LIKE '%s'" % database_name
    self._query(sql_command)
    result = self._fetch()
    return len(result) > 0

  def create_database(self, database_name):
    sql_command = "CREATE DATABASE %s" % database_name
    self._query(sql_command)


class AdminDbManager(DbManager):
  def __init__(self, database_name):
    db_user = DbUser.get_admin_user(database_name)
    if database_name is not None:
      db_user['database_name'] = DbUser.get_database_name(database_name)
    DbManager.__init__(self, db_user, "utf8")


class QuantDbManager(DbManager):
  def __init__(self, database_name=None):
    db_user = DbUser.get_readonly_user(database_name)
    if database_name is not None:
      db_user['database_name'] = DbUser.get_database_name(database_name)
    DbManager.__init__(self, db_user)
